import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import PIL
from PIL import ImageOps
import matplotlib.pyplot as plt

# 禁用GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# 加载图像
testimg = load_img('D:/project/dataset/segdata/images/Abyssinian_46.jpg', target_size=(160, 160))

# 转换为灰度图像
grayscale_img = testimg.convert('L')

# 将灰度图像转换为RGB格式以匹配模型输入要求
# 方法1: 直接将灰度图像转换为RGB（每个通道值相同）
rgb_from_gray = grayscale_img.convert('RGB')

# 安全加载模型，避免pydot相关错误
try:
    model = tf.keras.models.load_model('model.h5')
except Exception as e:
    print(f"模型加载出错: {e}")
    raise

# 预处理图像 - 使用灰度后转RGB的图像
img_array = img_to_array(rgb_from_gray)
img_array = img_array / 255.0  # 归一化
img_array = np.expand_dims(img_array, axis=0)  # 添加批次维度

# 执行预测
val_preds = model.predict(img_array)

def display_mask(i):
    """显示分割掩码"""
    mask = np.argmax(val_preds[i], axis=-1)
    # 维度调整
    mask = np.expand_dims(mask, axis=-1)
    img = PIL.ImageOps.autocontrast(tf.keras.preprocessing.image.array_to_img(mask))
    
    # 使用matplotlib显示图像
    plt.figure(figsize=(12, 6))
    
    # 显示原始图像
    plt.subplot(1, 3, 1)
    plt.imshow(testimg)
    plt.title('原始图像')
    plt.axis('off')
    
    # 显示灰度图像
    plt.subplot(1, 3, 2)
    plt.imshow(grayscale_img, cmap='gray')
    plt.title('灰度图像')
    plt.axis('off')
    
    # 显示分割结果
    plt.subplot(1, 3, 3)
    plt.imshow(img, cmap='gray')
    plt.title('分割结果')
    plt.axis('off')
    
    plt.tight_layout()
    plt.show()

# 显示分割结果
display_mask(0)

# 保存结果
result_img = PIL.ImageOps.autocontrast(tf.keras.preprocessing.image.array_to_img(
    np.expand_dims(np.argmax(val_preds[0], axis=-1), axis=-1)
))
result_img.convert('L').save('segmentation_result.png')
print("分割结果已保存为 segmentation_result.png")